import os
import requests
from langchain.prompts import PromptTemplate
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
import torch

class DocumentStorage:
    def __init__(self, faiss_index_path="faiss_index/index.faiss"):
        self.faiss_index_path = faiss_index_path
        self.embeddings = HuggingFaceEmbeddings(
            model_name="/mnt/langchain/model/bge-large-zh",
            model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
        )
        self.vector_store = None

    def load_faiss_vector_store(self):
        if os.path.exists(self.faiss_index_path):
            print(f"加载本地FAISS向量存储: {self.faiss_index_path}")
            self.vector_store = FAISS.load_local(
                self.faiss_index_path,
                self.embeddings,
                allow_dangerous_deserialization=True
            )
        else:
            print(f"未检测到向量存储文件 {self.faiss_index_path}")

    def search_similar_documents(self, query, top_k=3):
        if not self.vector_store:
            print("❌ FAISS 向量存储未加载，请先创建或加载存储")
            return []
        similar_docs = self.vector_store.similarity_search(query, k=top_k)
        print(f"查询 {query} 的相关内容:\n")
        for i, doc in enumerate(similar_docs):
            print(f"文档 {i + 1}:\n{doc.page_content[:200]}\n")
        return [doc.page_content for doc in similar_docs]


class APILLM:
    def __init__(self, api_url, api_key, model_name, temperature=0.3):
        self.api_url = api_url
        self.api_key = api_key
        self.model_name = model_name
        self.temperature = temperature

    def invoke(self, prompt):
        """调用阿里云百炼API生成响应"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        payload = {
            "model": self.model_name,
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": prompt
                    }
                ]
            },
            "parameters": {
                "temperature": self.temperature,
                "max_tokens": 1500
            }
        }

        try:
            response = requests.post(self.api_url, json=payload, headers=headers)
            response.raise_for_status()
            result = response.json()
            print("API 原始响应:", result)  # 调试信息

            # 处理认证错误
            if 'code' in result and result['code'] == 'InvalidApiKey':
                return "错误: API 密钥无效或未提供。"

            # 提取回答内容（根据实际结构修改）
            if 'output' in result and 'text' in result['output']:
                return result['output']['text']
            else:
                return "抱歉，无法解析回答。"
        except requests.exceptions.RequestException as e:
            print(f"API调用出错: {e}")
            return "抱歉，无法获取回答。"


class ReactAgentSystem:
    def __init__(self, document_storage, api_url, api_key):
        self.document_storage = document_storage
        self.llm = APILLM(
            api_url=api_url,
            api_key=api_key,
            model_name="qwen-plus",  # 或其他您在百炼平台开通的Qwen模型
            temperature=0.3
        )

    def create_react_agent(self, query, related_docs):
        react_prompt = PromptTemplate(
            template=(
                "以下是与问题相关的企业文档内容：\n{docs}\n\n"
                "请基于文档信息回答问题：\n问题：{query}\n"
            ),
            input_variables=["query", "docs"]
        )
        return react_prompt.format(query=query, docs="\n".join(related_docs))

    def generate_response(self, query):
        related_docs = self.document_storage.search_similar_documents(query)
        print(f"相关文档{related_docs}")
        prompt = self.create_react_agent(query, related_docs)
        response = self.llm.invoke(prompt)
        print("\nReact Agent回答：\n", response)
        return response


class ZeroShotAgentSystem:
    def __init__(self, api_url, api_key):
        self.llm = APILLM(
            api_url=api_url,
            api_key=api_key,
            model_name="qwen-plus",  # 或其他您在百炼平台开通的Qwen模型
            temperature=0.3
        )

    def create_zero_shot_prompt(self, query):
        prompt = PromptTemplate(
            template="请基于常识回答下述问题：\n问题：{query}\n",
            input_variables=["query"]
        )
        return prompt.format(query=query)

    def generate_response(self, query):
        prompt = self.create_zero_shot_prompt(query)
        response = self.llm.invoke(prompt)
        print("\nZero-shot Agent回答：\n", response)
        return response


class EnterpriseAgentSystem:
    def __init__(self, faiss_index_path, api_url, api_key):
        self.document_storage = DocumentStorage(faiss_index_path)
        self.document_storage.load_faiss_vector_store()
        self.react_agent = ReactAgentSystem(self.document_storage, api_url, api_key)
        self.zero_shot_agent = ZeroShotAgentSystem(api_url, api_key)

    def execute_query(self, query, user_react=True):
        if user_react:
            print("使用ReAct Agent进行回答：")
            response = self.react_agent.generate_response(query)
        else:
            print("使用Zero-shot Agent进行回答：")
            response = self.zero_shot_agent.generate_response(query)

        print(f"最终回答：{response}")
        return response


if __name__ == "__main__":
    # 初始化企业代理系统
    faiss_index_path = "faiss_index"

    # 阿里云百炼API配置
    api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
    # 请将下面的API_KEY替换为您在阿里云百炼平台获取的实际API Key
    api_key = "sk-973b48c9d559461bb46a60ebc340cad0"

    # 检查API Key是否已设置
    # if api_key == "sk-973b48c9d559461bb46a60ebc340cad0":
    #     print("请设置您的阿里云百炼API Key")
    #     print("可以通过环境变量DASHSCOPE_API_KEY设置，或直接修改代码中的api_key变量")
    #     exit(1)

    enterprise_agent = EnterpriseAgentSystem(faiss_index_path, api_url, api_key)

    print("欢迎使用企业问答系统！请输入您的问题（输入'退出'结束对话）")

    # 进入问答循环
    while True:
        # 获取用户输入
        user_query = input("\n请输入您的问题: ")

        # 检查是否退出
        if user_query.strip().lower() == "退出":
            print("感谢使用，再见！")
            break

        # 执行查询并获取结果
        print("查询结果:")
        enterprise_agent.execute_query(user_query, user_react=True)



